import unittest
import difflib
from AStock.ASUtil import ASHttpParser, date_to_bar_format
from AStock.ASData import financial_dict


class TestHttpParser(unittest.TestCase):

    def setUp(self):
        self._financial_dict = {item[3:]: item[:3] for item in financial_dict.keys()}

    def _fuzzy_match(self, cn_financial_item):
        matches = difflib.get_close_matches(
            cn_financial_item, self._financial_dict.keys(), n=1, cutoff=0.6)
        print('fuzzy match: {} -> {}'.format(cn_financial_item, matches))
        if matches:
            return self._financial_dict.get(matches[0])  # 只需要前三位数字代码
        return None

    def testParseTable(self):
        # url = 'http://stockdata.stock.hexun.com/2008/zcfz.aspx?stockid=600276&accountdate=2020.06.30'
        # url = 'http://vip.stock.finance.sina.com.cn/corp/go.php/vFD_FinancialGuideLine/stockid/600276/displaytype/4.phtml'
        url = 'https://money.finance.sina.com.cn/corp/go.php/vFD_FinancialGuideLine/stockid/600276/displaytype/4.phtml'
        header = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.125 Safari/537.36'
        }

        p = ASHttpParser(url, header=header)
        # df = p.parse_table(xpath='//div[@id="zaiyaocontent"]/table')
        df = p.parse_table(
            xpath='//table[@id="BalanceSheetNewTable0"]',
            df_handler=self.df_handler,
            report_date='20200331'
        )
        self.assertIsNotNone(df)
        print(df.head(30))
        for row in df.itertuples():
            # row[0], row[1], row[2] 等同于
            # row.Index, row.c1, row.c2
            field_name = self._fuzzy_match(row[1])
            if field_name is None:
                continue
            print('field name: {}, value: {}, type: {}'.format(field_name, row[2], type(row[2])))

    @staticmethod
    def df_handler(df, **kwarg):
        if df.shape[0] == 0:
            return df
        if 'report_date' in kwarg:
            report_date = date_to_bar_format(kwarg['report_date'])
            first_row = df.iloc[0]
            for index, value in first_row.items():
                if value == report_date:
                    # 取第0列和第index列
                    df = df.iloc[:, [0, index]]
        return df
